#!/usr/bin/python
# 书写时是习题后紧跟答案，不过排版时分离了习题与答案，把习题作为第一部分，答案作为第二部分

import sys
import os.path as op
import re
import logging


def get_match(s, index, begin, end):
    # 返回匹配的字符窜最后一个字符的指标下一个指标
    # 例如 { }; [ ]; \begin{env} \end{env}
    occurence = 0
    b_len, e_len = len(begin), len(end)
    assert s[index:index + len(begin)] == begin
    for i in range(index + len(begin), len(s)):
        if begin[0] != '\\' and s[i] == '\\':
            # begin 字符串不以 \ 开头时跳过转义字符
            i += 2
        elif i + b_len <= len(s) and s[i:i + b_len] == begin:
            # 容许嵌套
            occurence += 1
        elif i + e_len <= len(s) and s[i:i + e_len] == end:
            if occurence == 0:
                return i + e_len
            else:
                occurence -= 1


def get_exercises_solutions(s):
    # 内容基本被交错地划分为两种段：题目（及其他文字，包括章节名，行间文字），答案或注记
    # 期望之中题目与答案是一对，如果没有答案就加入空答案
    flags = {'exercise': 1, 'solution': -1, 'proof': -1, 'remark': 0}
    do_pair = 0
    i, exercises, solutions = 0, [], []

    def add_blank_solution(solutions):
        # 上一道习题没有答案时添加空答案
        nonlocal do_pair
        if do_pair not in [0, 1]:  # 用于排错
            logging.info(solutions)
            raise Exception(f'习题答案配对出错，do_pair 的值为 {do_pair}')
        if do_pair == 1:
            solutions.append('\n\\begin{solution}\n\\end{solution}\n')
            do_pair += flags['solution']

    while i < len(s):
        if s[i:].startswith('\\if'):
            # 跳过 \if ... \fi
            i = get_match(s, i, '\\if', '\\fi')
        elif s[i:].startswith('%'):
            # 跳过注释，这里未作判断，未考虑 \%
            i = get_match(s, i, '%', '\n')
        elif s[i:].startswith('\\section'):
            add_blank_solution(solutions)
            section_title_match = re.search(r'\\section\s*(\[.*\])?{[^}]*}', s[i:])
            exercises.append(section_title_match.group(0))
            solutions.append(section_title_match.group(0))
            i += section_title_match.end()
        elif s[i:].startswith('\\begin'):
            matched = re.search(r'^\\begin{([^}]*)}', s[i:])
            env = matched.group(1)
            k = get_match(s, i, f'\\begin{{{env}}}', f'\\end{{{env}}}')
            if env.startswith('exercise'):
                # 作业环境都是以 exercise 开头的
                # print(solutions)
                add_blank_solution(solutions)
                do_pair += flags['exercise']
                exercises.append(f'{s[i:k]}\n\n')
                solutions.append(s[i:k])  # 答案部分也包含进题目
            elif env in ['solution', 'proof', 'remark']:
                do_pair += flags[env]
                solutions.append(f'\n\n{s[i:k]}')
            else:
                exercises.append(s[i:k])
            i = k
        else:
            exercises.append(s[i])
            i += 1
    add_blank_solution(solutions)  # 最后一道题可能没有答案
    return ''.join(exercises), ''.join(solutions)


def collect(file, output):
    assert file.endswith('.tex')
    exercises, solutions = '', ''
    with open(file) as f:
        content = f.read()
    title_m = re.search(r'\\extitle{([0-9]*)}{([^}]*)}', content)
    chapter = title_m.group(1)
    title = title_m.group(2)
    this_exercises, this_solutions = get_exercises_solutions(content)
    exercises = '\n'.join([exercises, this_exercises, ''])
    solutions = '\n'.join([solutions, this_solutions, ''])
    content_new = '\n'.join([
        '\\setcounter{exercise}{0}\n',
        exercises,
        '\\printbibliography[title=参考文献]',
        f'\n\\newpage\n\\extitle[solution]{{{chapter}}}{{{title}}}',
        '\\setcounter{section}{0}\n',
        '\\setcounter{exercise}{0}\n',
        solutions,
    ])
    with open(output, 'w') as f:
        f.write(content_new)


if __name__ == '__main__':
    file = sys.argv[1]
    directory, filename = op.dirname(file), op.basename(file)
    output = sys.argv[2] if len(sys.argv) > 2 else op.join(directory, f'temp_{filename}')
    collect(file, output)
